import copy
import torch
import numpy as np
from sklearn.datasets import load_iris

#加载数据集
def load_data(shuffle=True):
    """
    加载鸢尾花数据
    输入：
        - shuffle：是否打乱数据，数据类型为bool
    输出：
        - X：特征数据，shape=[150,4]
        - y：标签数据, shape=[150,3]
    """
    #加载原始数据
    X = np.array(load_iris().data, dtype=np.float32)
    y = np.array(load_iris().target, dtype=np.int64)
    X = torch.tensor(X)
    y = torch.tensor(y)

    #数据归一化
    X_min = torch.min(X, dim=0)
    X_max = torch.max(X, dim=0)
    X = (X-X_min.values) / (X_max.values-X_min.values)

    #如果shuffle为True，随机打乱数据
    if shuffle:
        idx = torch.randperm(X.shape[0])
        X_new = copy.deepcopy(X)
        y_new = copy.deepcopy(y)
        for i in range(X.shape[0]):
            X_new[i] = X[idx[i]]
            y_new[i] = y[idx[i]]
        X = X_new
        y = y_new

    return X, y

